import React, { useState, useEffect, useRef } from 'react';
import { Input, Button, List, message, Card, Space, Select, Typography, Tabs } from 'antd';
import { SendOutlined, LinkOutlined, DisconnectOutlined } from '@ant-design/icons';
import './App.css';

const { TabPane } = Tabs;
const { Text } = Typography;

function App() {
  const [messages, setMessages] = useState({}); // 按房间存储消息
  const [inputValue, setInputValue] = useState('');
  const [connected, setConnected] = useState(false);
  const [loading, setLoading] = useState(false);
  const [currentRoom, setCurrentRoom] = useState('room1');
  const [rooms, setRooms] = useState(['room1', 'room2']); // 可用聊天室列表
  const ws = useRef(null);

  // 初始化消息存储
  useEffect(() => {
    const initialMessages = {};
    rooms.forEach(room => {
      initialMessages[room] = [];
    });
    setMessages(initialMessages);
  }, []);

  // 连接WebSocket
  const connectWebSocket = () => {
    setLoading(true);
    ws.current = new WebSocket(`ws://localhost:8000/ws/${currentRoom}`);
    
    ws.current.onopen = () => {
      setConnected(true);
      setLoading(false);
      message.success(`已连接到聊天室 ${currentRoom}`);
    };
    
    ws.current.onmessage = (e) => {
      try {
        const data = JSON.parse(e.data);
        setMessages(prev => ({
          ...prev,
          [currentRoom]: [...prev[currentRoom], data]
        }));
      } catch (err) {
        console.error('消息解析错误:', err);
      }
    };
    
    ws.current.onclose = () => {
      setConnected(false);
      message.warning('已断开WebSocket连接');
    };
    
    ws.current.onerror = (err) => {
      console.error('WebSocket错误:', err);
      message.error('连接错误');
      setLoading(false);
    };
  };

  // 断开连接
  const disconnectWebSocket = () => {
    if (ws.current) {
      ws.current.close();
    }
  };

  // 发送消息
  const sendMessage = () => {
    if (!connected) {
      message.warning('请先连接WebSocket');
      return;
    }
    
    if (!inputValue.trim()) {
      message.warning('消息不能为空');
      return;
    }
    
    const messageData = {
      content: inputValue,
      timestamp: new Date().getTime()
    };
    
    ws.current.send(JSON.stringify(messageData));
    setInputValue('');
  };

  // 切换房间
  const changeRoom = (room) => {
    setCurrentRoom(room);
    if (connected) {
      disconnectWebSocket();
      connectWebSocket();
    }
  };

  // 组件卸载时断开连接
  useEffect(() => {
    return () => {
      if (ws.current) {
        ws.current.close();
      }
    };
  }, []);

  return (
    <div className="app-container">
      <Card 
        title="多聊天室WebSocket示例" 
        bordered={false}
        extra={
          <Space>
            <Select
              value={currentRoom}
              style={{ width: 120 }}
              onChange={changeRoom}
              disabled={connected}
            >
              {rooms.map(room => (
                <Select.Option key={room} value={room}>
                  {room}
                </Select.Option>
              ))}
            </Select>
            
            {!connected ? (
              <Button
                type="primary"
                icon={<LinkOutlined />}
                onClick={connectWebSocket}
                loading={loading}
              >
                连接
              </Button>
            ) : (
              <Button
                danger
                icon={<DisconnectOutlined />}
                onClick={disconnectWebSocket}
              >
                断开
              </Button>
            )}
          </Space>
        }
      >
        <Tabs activeKey={currentRoom} onChange={changeRoom}>
          {rooms.map(room => (
            <TabPane tab={room} key={room}>
              <div className="chat-container">
                <List
                  className="message-list"
                  itemLayout="horizontal"
                  dataSource={messages[room] || []}
                  renderItem={item => (
                    <List.Item>
                      <List.Item.Meta
                        title={`${item.client} - ${item.time}`}
                        description={item.content}
                      />
                    </List.Item>
                  )}
                />
                
                <div className="input-area">
                  <Input
                    value={inputValue}
                    onChange={e => setInputValue(e.target.value)}
                    onPressEnter={sendMessage}
                    placeholder="输入消息..."
                    disabled={!connected}
                  />
                  <Button
                    type="primary"
                    icon={<SendOutlined />}
                    onClick={sendMessage}
                    disabled={!connected}
                  >
                    发送
                  </Button>
                </div>
              </div>
            </TabPane>
          ))}
        </Tabs>
      </Card>
    </div>
  );
}

export default App;